# register_speakers.py
import numpy as np
import torch
from speechbrain.pretrained import EncoderClassifier
import torchaudio

# 加载声纹模型（首次运行会下载到本地，之后离线可用）
classifier = EncoderClassifier.from_hparams(
    source="speechbrain/spkrec-ecapa-voxceleb",
    savedir="local_models/ecapa-voxceleb",
    run_opts={"device": "cuda"}  # 或 "cpu"
)

def extract_embedding(audio_path):
    signal, sr = torchaudio.load(audio_path)
    if sr != 16000:
        signal = torchaudio.transforms.Resample(sr, 16000)(signal)
    with torch.no_grad():
        emb = classifier.encode_batch(signal)
    return emb.squeeze().cpu().numpy()

# 注册用户
speaker_db = {}
speaker_db["张三"] = extract_embedding("samples/zhangsan.wav")
speaker_db["李四"] = extract_embedding("samples/lisi.wav")

# 保存到本地
np.save("speaker_db.npy", speaker_db)
print("声纹库已保存至 speaker_db.npy")